{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Multiobjective optimization via functional operators API\n",
    "\n",
    "The functional operators API of EvoTorch (`evotorch.operators.functional`) can be used for multiobjective optimization.\n",
    "In this notebook, we demonstrate how this functional operators API can be used to tackle the Kursawe function, which has two objectives to be minimized."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "We begin with the necessary imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import evotorch.operators.functional as func_ops\n",
    "from evotorch.decorators import rowwise\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Below, we implement Kursawe's function.\n",
    "\n",
    "Notice how we decorate the function via `evotorch.decorators.rowwise`. This `@rowwise` decorator allows us to implement the function `f` with the simple assumption that `x` is a single vector. However, when calling this decorated function `f` from outside, we will be able to provide `x` as a matrix, in which case the `@rowwise` decorator will broadcast the function `f` such that it will be applied for each row of the matrix. In fact, we can even give a tensor with 3 or more dimensions as `x`, and the decorated `f` will interpret all of the extra leftmost dimensions as the batch dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@rowwise\n",
    "def f(x: torch.Tensor) -> torch.Tensor:\n",
    "    # Kursawe's function\n",
    "\n",
    "    f1 = torch.sum(\n",
    "        -10 * torch.exp(\n",
    "            -0.2 * torch.sqrt(x[0:2] ** 2.0 + x[1:3] ** 2.0)\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    f2 = torch.sum(\n",
    "        (torch.abs(x) ** 0.8) + (5 * torch.sin(x ** 3)),\n",
    "    )\n",
    "    fitnesses = torch.hstack([f1, f2])\n",
    "    return fitnesses"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "Below, we have the constants regarding the problem, and hyperparameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\"\n",
    "solution_length = 3\n",
    "objective_sense = [\"min\", \"min\"]\n",
    "lb = -5.0\n",
    "ub = 5.0\n",
    "\n",
    "popsize = 200\n",
    "num_generations = 100\n",
    "mutation_stdev = 0.03\n",
    "tournament_size = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Initialize a population, and store it via the variable `population`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "population = (torch.rand(popsize, solution_length, device=device) * (ub - lb)) + lb\n",
    "population.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Evaluate the initial population, and store the evaluation results within the variable `evals`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "evals = f(population)\n",
    "evals.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "Main loop of the optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "for generation in range(1, 1 + num_generations):\n",
    "\n",
    "    # Apply a tournament selection, and a simulated binary cross-over (SBX)\n",
    "    # on the selected parents\n",
    "    candidates = func_ops.simulated_binary_cross_over(\n",
    "        population,\n",
    "        evals,\n",
    "        tournament_size=tournament_size,\n",
    "        eta=1,\n",
    "        objective_sense=objective_sense,\n",
    "    )\n",
    "\n",
    "    # Instead of a simulated binary cross-over, we could also use a two-point\n",
    "    # cross-over, as follows:\n",
    "    #\n",
    "    # candidates = func_ops.two_point_cross_over(\n",
    "    #     population,\n",
    "    #     evals,\n",
    "    #     tournament_size=tournament_size,\n",
    "    #     objective_sense=objective_sense,\n",
    "    # )\n",
    "\n",
    "    # Apply Gaussian mutation on the results of the cross-over operation\n",
    "    candidates = candidates + (torch.randn_like(candidates) * mutation_stdev)\n",
    "\n",
    "    # Evaluate the mutated candidate solutions\n",
    "    candidate_evals = f(candidates)\n",
    "\n",
    "    # Form an extended population by combining the parent solutions and the\n",
    "    # candidate solutions.\n",
    "    extended_population, extended_evals = func_ops.combine(\n",
    "        (population, evals),\n",
    "        (candidates, candidate_evals),\n",
    "        # We are passing `objective_sense` to inform the `combine` function\n",
    "        # that the problem at hand is multi-objective:\n",
    "        objective_sense=objective_sense,\n",
    "    )\n",
    "\n",
    "    # Take the `popsize` number of solutions from best pareto-fronts.\n",
    "    population, evals = func_ops.take_best(\n",
    "        extended_population,\n",
    "        extended_evals,\n",
    "        popsize,\n",
    "        objective_sense=objective_sense,\n",
    "        # When selecting the solutions, we want the crowding distances of the\n",
    "        # solutions to be taken into account:\n",
    "        crowdsort=True\n",
    "    )\n",
    "\n",
    "    # Print the current status:\n",
    "    print(\"Generation:\", generation, \"  Best evals of the population:\", torch.max(evals, dim=0).values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "Considering that `evals` now stores the evaluation results of the latest population, we can take the best solutions belonging to the best pareto-front as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute domination count (i.e. how many times a solution was dominated)\n",
    "# for each solution\n",
    "dcounts = func_ops.domination_counts(evals, objective_sense=objective_sense)\n",
    "\n",
    "# Make a mask in which the i-th element is True if the i-th solution of the\n",
    "# population has never been dominated\n",
    "# (i.e. if the i-th solution is at the best pareto-front)\n",
    "at_best_front = (dcounts == 0)\n",
    "\n",
    "# Filter both the decision values tensor and the evaluation results tensor\n",
    "# such that only the solutions on the best pareto-front will be included.\n",
    "# The results of this filtering operation will be stored by the variables\n",
    "# `best_pop` and `best_evals`.\n",
    "best_pop = population[at_best_front]\n",
    "best_evals = evals[at_best_front]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_pop.shape, best_evals.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "Plot the fitnesses of the best solutions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(f\"Fitnesses after {num_generations} generations\")\n",
    "plt.scatter(best_evals[:, 0].cpu().numpy(), best_evals[:, 1].cpu().numpy())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Batched multiobjective optimization\n",
    "\n",
    "The functional operators API of EvoTorch is written in such a way that a single call to an operator can work on not just a single population, but on a batch of multiple populations, in a vectorized manner.\n",
    "\n",
    "Below, we demonstrate this feature by modifying the multiobjective example above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us consider 4 populations:\n",
    "num_populations = 4\n",
    "\n",
    "# Size for each population:\n",
    "popsize = 200\n",
    "\n",
    "# Shared hyperparameters\n",
    "num_generations = 30\n",
    "tournament_size = 4\n",
    "mutation_stdev = 0.03\n",
    "\n",
    "# Hyperparameters that vary for each population:\n",
    "eta = torch.tensor([1.0, 8.0, 20.0, 40.0], device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a batch of populations (each population is initialized the same)\n",
    "population = (torch.rand(popsize, solution_length, device=device) * (ub - lb)) + lb\n",
    "evals = f(population)\n",
    "broadcaster = torch.ones(num_populations, 1, 1, device=device)\n",
    "population = population * broadcaster\n",
    "evals = evals * broadcaster\n",
    "\n",
    "# Alternatively, in some cases, you might want to do the initialization in\n",
    "# such a way that each population within the batch is different:\n",
    "#\n",
    "# population = (torch.rand(num_populations, popsize, solution_length, device=device) * (ub - lb)) + lb\n",
    "# evals = f(population)\n",
    "\n",
    "population.shape, evals.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "for generation in range(1, 1 + num_generations):\n",
    "    candidates = func_ops.simulated_binary_cross_over(\n",
    "        population,\n",
    "        evals,\n",
    "        tournament_size=tournament_size,\n",
    "        objective_sense=objective_sense,\n",
    "        #\n",
    "        # Upon seeing that `eta` is given as a vector (instead of a scalar),\n",
    "        # the function `simulated_binary_cross_over` will treat the `eta` as\n",
    "        # a batch of hyperparameters. In more details, the first `eta` is\n",
    "        # used on the first population of the batch, the second `eta` is used\n",
    "        # on the second population of the batch, and so on...\n",
    "        eta=eta,\n",
    "    )\n",
    "\n",
    "    candidates = candidates + (torch.randn_like(candidates) * mutation_stdev)\n",
    "    candidate_evals = f(candidates)\n",
    "\n",
    "    extended_population, extended_evals = func_ops.combine(\n",
    "        (population, evals),\n",
    "        (candidates, candidate_evals),\n",
    "        objective_sense=objective_sense,\n",
    "    )\n",
    "\n",
    "    population, evals = func_ops.take_best(\n",
    "        extended_population,\n",
    "        extended_evals,\n",
    "        popsize,\n",
    "        objective_sense=objective_sense,\n",
    "        crowdsort=True\n",
    "    )\n",
    "\n",
    "    # Print the current status:\n",
    "    print(\"Generation:\", generation)\n",
    "    print(\"Best evals of the populations:\")\n",
    "    print(torch.max(evals, dim=1).values)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "For each solution within each population, compute the domination count:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "dcounts = func_ops.domination_counts(evals, objective_sense=objective_sense)\n",
    "dcounts.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "From each population, take the best pareto-front, and plot the fitnesses belonging to that pareto-front:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i_population in range(num_populations):\n",
    "    single_pop_dcounts = dcounts[i_population, :]\n",
    "    single_pop = population[i_population, :, :]\n",
    "    single_pop_evals = evals[i_population, :, :]\n",
    "\n",
    "    at_best_front = (single_pop_dcounts == 0)\n",
    "\n",
    "    best_pop = single_pop[at_best_front]\n",
    "    best_evals = single_pop_evals[at_best_front]\n",
    "\n",
    "    plt.title(f\"Fitnesses with eta={float(eta[i_population].cpu())}, after {num_generations} generations\")\n",
    "    plt.scatter(best_evals[:, 0].cpu().numpy(), best_evals[:, 1].cpu().numpy())\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
